import json
import matplotlib
matplotlib.use('Agg') 
import matplotlib.pyplot as plt
import numpy as np
import os

# 加载训练数据
def load_train_data(json_path):
    with open(json_path, 'r') as f:
        data = json.load(f)
    return data

# 可视化损失曲线
def visualize_losses(weight_path, train_data):
    # 提取数据
    epochs = [entry['Epoch'] for entry in train_data]
    loss_G = [entry['avg_loss_G'] for entry in train_data]
    loss_D = [entry['avg_loss_D'] for entry in train_data]
    
    # 设置全局绘图参数
    plt.rcParams.update({
        'figure.figsize': (12, 6),
        'font.size': 14,
        'axes.grid': True,
        'grid.linestyle': '--',
        'grid.alpha': 0.6
    })
    
    # 创建生成器损失图
    plt.figure()
    plt.plot(epochs, loss_G, 'b-o', linewidth=2, markersize=8, label='Generator Loss')
    
    # 标记损失最小值
    min_idx = np.argmin(loss_G)
    plt.scatter(epochs[min_idx], loss_G[min_idx], c='red', s=150, 
                zorder=5, label=f'Min Loss: {loss_G[min_idx]:.4f} (Epoch {epochs[min_idx]})')
    
    # 设置图形参数
    plt.title('Generator Loss During Training', fontsize=18, pad=20)
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    # plt.ylim(1, 6)
    plt.legend(loc='upper right')
    plt.tight_layout()
    
    # 保存生成器损失图
    plt.savefig(os.path.join(weight_path, "generator_loss.png"), dpi=300)
    print("已保存生成器损失图: generator_loss.png")
    
    # 创建判别器损失图
    plt.figure()
    plt.plot(epochs, loss_D, 'g-s', linewidth=2, markersize=8, label='Discriminator Loss')
    
    # 添加趋势线（二次多项式拟合）
    z = np.polyfit(epochs, loss_D, 2)
    p = np.poly1d(z)
    trend = p(epochs)
    plt.plot(epochs, trend, 'r--', linewidth=2.5, alpha=0.7, 
             label='Trend Line (Quadratic Fit)')
    
    # 设置图形参数
    plt.title('Discriminator Loss During Training', fontsize=18, pad=20)
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    # plt.ylim(0, 0.025)
    plt.legend(loc='upper right')
    plt.tight_layout()
    
    # 保存判别器损失图
    plt.savefig(os.path.join(weight_path, "discriminator_loss.png"), dpi=300)
    print("已保存判别器损失图: discriminator_loss.png")


# 主函数
if __name__ == "__main__":

    import argparse 
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--weight_path", 
        type=str, default=None, 
        help="Path to save weight (e.g., 'UnetCBAM/weights/2025-09-10-20-47-37')"
    )
    args = parser.parse_args()

    # 替换为你的训练数据JSON文件路径
    weight_path = args.weight_path


    JSON_PATH = os.path.join(weight_path, "train_data.json")

    # 加载并可视化
    train_data = load_train_data(JSON_PATH)
    visualize_losses(weight_path, train_data)